﻿from sklearn import datasets
import matplotlib.pyplot as plt
import numpy as np

iris = datasets.load_iris()
irisFeatures = iris["data"]
irisFeaturesName = iris["feature_names"]
irisLabels = iris["target"]



class KMeans:
    def __init__(self, k: int, n: int):
        self.K = k
        self.N = n
        self.u = np.zeros((k, n))
        self.C = [[] for _ in range(k)]

    def select_u0(self, data: np.ndarray):
        indices = np.random.choice(data.shape[0], self.K, replace=False)
        self.u = data[indices]

    def fit(self, data: np.ndarray):
        self.select_u0(data)
        J = float('inf')

        cnt=0
        oldJ=[]

        while True:
            cnt+=1

            self.C = [[] for _ in range(self.K)]
            distances = np.linalg.norm(data[:, np.newaxis] - self.u, axis=2)
            c_centroids = np.argmin(distances, axis=1)

            for i in range(self.K):
                self.C[i] = data[c_centroids == i]
            new_centroids = np.array([np.mean(self.C[i], axis=0) if len(self.C[i]) > 0 else self.u[i] for i in range(self.K)])
            new_J = np.sum(np.min(distances, axis=1))

            oldJ.append(float(abs(new_J-J)))

            if abs(new_J - J) < 0.001:
                break
            J = new_J
            self.u = new_centroids

    def random_select_u0(self,data:np.ndarray):
        self.u[0] = data[np.random.randint(len(data))]
        for i in range(1, self.K):
            distances = np.array([min([np.linalg.norm(x - u) ** 2 for u in self.u[:i]]) for x in data])
            probabilities = distances / np.sum(distances)
            cumulative_probabilities = np.cumsum(probabilities)
            rand = np.random.rand()
            index = np.where(cumulative_probabilities >= rand)[0][0]
            self.u[i] = data[index]

    def draw(self,cnt,oldJ):
        plt.plot(range(cnt),oldJ,"r-")
        plt.xlabel("iteration")
        plt.ylabel("abs(J-oldJ)")
        plt.show()

k=6
c=7
c=min(k,c)

model = KMeans(k, 4)
model.fit(irisFeatures)
colors = ['red', 'green', 'blue', 'purple', 'magenta', 'cyan', 'yellow','orange']
markers = ['o', '*', '+', '1', '2', '3', '4','5']

for i, color, marker in zip(range(c), colors, markers):
    if len(model.C[i]) > 0:
        x = np.array(model.C[i])
        plt.scatter(x[:, 0], x[:, 1], c=color, marker=marker, label=f'cluster{i+1}')

plt.scatter(model.u[:, 0], model.u[:, 1], c="black", marker='x', label='centers')

plt.xlabel('petal length')
plt.ylabel('petal width')
plt.legend(loc='best')
plt.title('KMeans')
plt.show()

# def _draw(k):
#     c = 8
#     c = min(k, c)
#
#     model = KMeans(k, 4)
#     model.fit(irisFeatures)
#     colors = ['red', 'green', 'blue', 'purple', 'magenta', 'cyan', 'yellow', 'orange']
#     markers = ['o', '*', '+', '1', '2', '3', '4', '5']
#
#     for i, color, marker in zip(range(c), colors, markers):
#         if len(model.C[i]) > 0:
#             x = np.array(model.C[i])
#             plt.scatter(x[:, 0], x[:, 1], c=color, marker=marker, label=f'cluster{i + 1}')
#
#     plt.scatter(model.u[:, 0], model.u[:, 1], c="black", marker='x', label='centers')
#
#     plt.xlabel('petal length')
#     plt.ylabel('petal width')
#     plt.legend(loc='best')
#     plt.title('KMeans')
#     plt.show()
#
# for i in range(7):
#     _draw(i+1)